#!/usr/bin/env python3
import argparse
import os
import re
import shutil
import struct
import subprocess
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from elftools.elf.elffile import ELFFile


@dataclass
class TargetSpec:
    name: str
    max_params: int


@dataclass
class ParamHit:
    index: int
    reg: str
    value: Optional[int]
    string: Optional[str]


@dataclass
class CallHit:
    target: str
    callsite: int
    params: List[ParamHit]
    caller: Optional[str] = None


MAX_REGISTER_PARAMS = 8


def run_objdump_disasm(path: str) -> List[str]:
    """Disassemble .text with symbols and reloc info using objdump.

    Returns objdump output lines. Requires `objdump` (binutils).
    """
    # Prefer cross-objdump if available
    candidates = [
        "aarch64-linux-gnu-objdump",
        "objdump",
    ]
    ob = None
    for c in candidates:
        if shutil.which(c):
            ob = c
            break
    if not ob:
        raise RuntimeError("No objdump found. Install binutils or binutils-aarch64-linux-gnu.")

    # Use raw register names to simplify parsing, include relocs, demangle off.
    # We disassemble all sections to catch stubs if .text split.
    cmd = [ob, "-d", "-r", "-w", "-M", "reg-names-raw", path]
    try:
        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        # Retry without reg-names-raw if unsupported
        try:
            out = subprocess.check_output([ob, "-d", "-r", "-w", path], stderr=subprocess.STDOUT)
            return out.decode(errors="ignore").splitlines()
        except subprocess.CalledProcessError:
            pass
        # Retry with -m aarch64 if using host objdump that supports multi-arch
        if ob == "objdump":
            try:
                out = subprocess.check_output([ob, "-m", "aarch64", "-d", "-r", "-w", "-M", "reg-names-raw", path], stderr=subprocess.STDOUT)
                return out.decode(errors="ignore").splitlines()
            except subprocess.CalledProcessError as e2:
                msg = e2.output.decode(errors='ignore')
                raise RuntimeError("objdump failed. Install aarch64 objdump (binutils-aarch64-linux-gnu).\n" + msg)
        raise RuntimeError(f"objdump failed: {e.output.decode(errors='ignore')}")
    return out.decode(errors="ignore").splitlines()


def load_elf(path: str) -> ELFFile:
    f = open(path, "rb")
    return ELFFile(f)


def va_to_file_off(elf: ELFFile, va: int) -> Optional[int]:
    # Map virtual address to file offset using PT_LOAD segments
    for seg in elf.iter_segments():
        if seg["p_type"] != "PT_LOAD":
            continue
        vaddr = seg["p_vaddr"]
        memsz = seg["p_memsz"]
        filesz = seg["p_filesz"]
        off = seg["p_offset"]
        if va >= vaddr and va < vaddr + memsz:
            rel = va - vaddr
            if rel < filesz:
                return off + rel
            else:
                # Not in file-backed range
                return None
    return None


def read_c_string_at_va(elf: ELFFile, va: int, max_len: int = 4096) -> Optional[str]:
    off = va_to_file_off(elf, va)
    if off is None:
        return None
    stream = elf.stream
    stream.seek(off)
    data = bytearray()
    for _ in range(max_len):
        b = stream.read(1)
        if not b:
            break
        if b == b"\x00":
            break
        data.extend(b)
    try:
        return data.decode("utf-8")
    except UnicodeDecodeError:
        try:
            return data.decode("latin-1")
        except Exception:
            return None


def normalize_symbol_name(name: str) -> str:
    base = name
    if base.endswith('@plt'):
        base = base[:-4]
    if '@@' in base:
        base = base.split('@@', 1)[0]
    if '@' in base:
        base = base.split('@', 1)[0]
    if base.endswith('.plt'):
        base = base[:-4]
    return base


_FUNC_HDR_RE = re.compile(r"^([0-9a-fA-F]+) <([^>]+)>:")
_INST_RE = re.compile(r"^\s*([0-9a-fA-F]+):\s+([0-9a-fA-F ]+)\s+\t(.*)$")


def parse_objdump(lines: List[str]) -> List[Tuple[str, int, int, str, Optional[int]]]:
    """Parse objdump output and yield (func, addr, line_no, asm, insn32) for each instruction."""
    cur_func = "?"
    result: List[Tuple[str, int, int, str, Optional[int]]] = []
    for i, line in enumerate(lines):
        m = _FUNC_HDR_RE.match(line)
        if m:
            cur_func = m.group(2)
            continue
        mi = _INST_RE.match(line)
        if not mi:
            continue
        addr = int(mi.group(1), 16)
        bytes_str = mi.group(2).strip().replace(" ", "")
        insn32: Optional[int] = None
        if len(bytes_str) >= 8:
            try:
                # Take the last 8 hex chars (one 32-bit instruction)
                insn32 = int(bytes_str[-8:], 16)
            except ValueError:
                insn32 = None
        asm = mi.group(3)
        result.append((cur_func, addr, i, asm, insn32))
    return result


def extract_call_target(asm: str) -> Optional[str]:
    s = asm.strip()
    if not s.startswith('bl'):
        return None
    parts = s.split(None, 1)
    if len(parts) < 2:
        return None
    tail = parts[1]
    m = re.search(r'<([^>]+)>', tail)
    if m:
        return m.group(1)
    token = tail.split()[0]
    return token


def try_extract_reg_write(asm: str) -> Optional[Tuple[str, int]]:
    """Recognize patterns that load absolute addresses into a register.

    Supported AArch64 idioms:
    - adrp xN, SYMBOL; add xM, xN, #imm (pair handled in tracker)
    - adrp xN, SYMBOL; ldr xM, [xN, #imm]
    - adrp/add to same register x0
    - movz/movk sequence to synthesize 64-bit immediate
    Returns (reg, value) when the instruction alone resolves to an absolute immediate, else None.
    """
    s = asm.strip()
    # movz/movk pattern
    m = re.match(r"movz\s+([xw]\d+),\s*#0x([0-9a-fA-F]+)(?:,\s*lsl\s*#(\d+))?", s)
    if m:
        reg = m.group(1)
        if reg.startswith('w'):
            reg = 'x' + reg[1:]
        imm = int(m.group(2), 16)
        shift = int(m.group(3) or 0)
        return reg, imm << shift
    m = re.match(r"movk\s+([xw]\d+),\s*#0x([0-9a-fA-F]+)(?:,\s*lsl\s*#(\d+))?", s)
    if m:
        reg = m.group(1)
        if reg.startswith('w'):
            reg = 'x' + reg[1:]
        imm = int(m.group(2), 16)
        shift = int(m.group(3) or 0)
        return reg, (imm << shift) | 0  # handled by tracker merge
    # Direct move from another reg or immediate (rare immediate)
    m = re.match(r"mov\s+([xw]\d+),\s*#?0x([0-9a-fA-F]+)", s)
    if m:
        reg = m.group(1)
        if reg.startswith('w'):
            reg = 'x' + reg[1:]
        imm = int(m.group(2), 16)
        return reg, imm
    return None


def scan_target_calls(instrs: List[Tuple[str, int, int, str, Optional[int]]], target_names: List[str]) -> List[Tuple[int, str, int, str]]:
    wanted = set(target_names)
    calls: List[Tuple[int, str, int, str]] = []
    for idx, (func, addr, lno, asm, _) in enumerate(instrs):
        target = extract_call_target(asm)
        if not target:
            continue
        base = normalize_symbol_name(target)
        if base in wanted:
            calls.append((idx, func, addr, base))
    return calls


def resolve_regs_before_call(elf: ELFFile, instrs: List[Tuple[str, int, int, str, Optional[int]]], call_index: int, target_regs: List[str]) -> Dict[str, Optional[int]]:
    """Walk backwards from call_index to resolve a list of registers as absolute values.

    Supports common AArch64 idioms (adrp/add, adrp/ldr, movz/movk, stack temps).
    Returns mapping of requested register names to resolved virtual addresses (if any).
    """
    reg_val: Dict[str, int] = {}
    reg_page_sym: Dict[str, str] = {}
    reg_page_addr: Dict[str, int] = {}
    reg_parts: Dict[str, Dict[int, int]] = {}
    stack_slots: Dict[int, int] = {}
    pending_sp_loads: List[Tuple[str, int]] = []
    pending_sp_store: Dict[int, str] = {}
    pending_add_imm: Dict[str, int] = {}
    pending_add_dep: Dict[str, Tuple[str, int]] = {}
    targets = set(target_regs)

    def fulfill_with_reg(reg: str) -> None:
        if reg in reg_val:
            for dst, (base, imm) in list(pending_add_dep.items()):
                if base == reg:
                    reg_val[dst] = reg_val[base] + imm
                    del pending_add_dep[dst]
                    fulfill_with_reg(dst)
        for off, sreg in list(pending_sp_store.items()):
            if sreg == reg and reg in reg_val:
                stack_slots[off] = reg_val[reg]
                for (dst2, off2) in list(pending_sp_loads):
                    if off2 == off:
                        reg_val[dst2] = stack_slots[off]
                        pending_sp_loads.remove((dst2, off2))
                del pending_sp_store[off]

    def done() -> bool:
        return all(reg in reg_val for reg in targets) and not pending_sp_loads and not pending_add_dep

    start = max(0, call_index - 80)
    for idx in range(call_index - 1, start - 1, -1):
        if done():
            break
        _, pc, _, asm, insn32 = instrs[idx]
        s = asm.strip()

        if s.startswith('ret') or s.startswith('blr ') or s.startswith('bl '):
            break

        m = re.match(r"adrp\s+([xw]\d+),\s*([._A-Za-z0-9$@:+\-/]+)", s)
        if m:
            reg = m.group(1)
            if reg.startswith('w'):
                reg = 'x' + reg[1:]
            sym = strip_lo12_suffix(m.group(2))
            reg_page_sym[reg] = sym
            if insn32 is not None:
                base = decode_adrp_page_base(insn32, pc)
                if base is not None:
                    reg_page_addr[reg] = base
                    if reg in pending_add_imm:
                        reg_val[reg] = base + pending_add_imm[reg]
                        fulfill_with_reg(reg)
            continue

        m = re.match(r"add\s+([xw]\d+),\s*([xw]\d+),\s*(?:#0x([0-9a-fA-F]+)|([._A-Za-z0-9$@:+\-/]+))", s)
        if m:
            dst, base = m.group(1), m.group(2)
            if dst.startswith('w'):
                dst = 'x' + dst[1:]
            if base.startswith('w'):
                base = 'x' + base[1:]
            if m.group(3):
                imm = int(m.group(3), 16)
                if base in reg_val:
                    reg_val[dst] = reg_val[base] + imm
                    fulfill_with_reg(dst)
                elif base in reg_page_addr:
                    reg_val[dst] = reg_page_addr[base] + imm
                    fulfill_with_reg(dst)
                else:
                    if dst == base:
                        pending_add_imm[dst] = imm
                    else:
                        pending_add_dep[dst] = (base, imm)
            else:
                sym = strip_lo12_suffix(m.group(4))
                if base in reg_page_sym and reg_page_sym[base] == sym:
                    addr = lookup_symbol_va(elf, sym)
                    if addr is not None:
                        reg_val[dst] = addr
                        fulfill_with_reg(dst)
            continue

        m = re.match(r"ldr\s+([xw]\d+),\s*\[([xw]\d+),\s*#(0x[0-9a-fA-F]+|[0-9]+)\]", s)
        if m:
            dst, base = m.group(1), m.group(2)
            if dst.startswith('w'):
                dst = 'x' + dst[1:]
            if base.startswith('w'):
                base = 'x' + base[1:]
            imm = int(m.group(3), 0)
            if base in reg_val:
                reg_val[dst] = reg_val[base] + imm
                fulfill_with_reg(dst)
            elif base == 'sp':
                if imm in stack_slots:
                    reg_val[dst] = stack_slots[imm]
                else:
                    pending_sp_loads.append((dst, imm))
            continue

        m = re.match(r"str\s+([xw]\d+),\s*\[sp,\s*#(0x[0-9a-fA-F]+|[0-9]+)\]", s)
        if m:
            src = m.group(1)
            if src.startswith('w'):
                src = 'x' + src[1:]
            imm = int(m.group(2), 0)
            if src in reg_val:
                stack_slots[imm] = reg_val[src]
                for (dst2, off2) in list(pending_sp_loads):
                    if off2 == imm:
                        reg_val[dst2] = stack_slots[imm]
                        pending_sp_loads.remove((dst2, off2))
            else:
                pending_sp_store[imm] = src
            continue

        m = re.match(r"adr\s+([xw]\d+),\s*([._A-Za-z0-9$@:+\-/]+)", s)
        if m:
            reg = m.group(1)
            if reg.startswith('w'):
                reg = 'x' + reg[1:]
            sym = strip_lo12_suffix(m.group(2))
            addr = lookup_symbol_va(elf, sym)
            if addr is not None:
                reg_val[reg] = addr
                fulfill_with_reg(reg)
            continue

        extracted = try_extract_reg_write(s)
        if extracted:
            reg, val = extracted
            if reg.startswith('w'):
                reg = 'x' + reg[1:]
            if s.startswith('movz'):
                reg_parts[reg] = {0: 0, 16: 0, 32: 0, 48: 0}
                match_shift = re.search(r"lsl\s*#(\d+)", s)
                shift = int(match_shift.group(1)) if match_shift else 0
                reg_parts[reg][shift] = (val >> shift) & 0xFFFF
                reg_val[reg] = combine_mov_parts(reg_parts[reg])
            elif s.startswith('movk'):
                if reg not in reg_parts:
                    reg_parts[reg] = {0: 0, 16: 0, 32: 0, 48: 0}
                match_shift = re.search(r"lsl\s*#(\d+)", s)
                shift = int(match_shift.group(1)) if match_shift else 0
                reg_parts[reg][shift] = (val >> shift) & 0xFFFF
                reg_val[reg] = combine_mov_parts(reg_parts[reg])
            else:
                reg_val[reg] = val
            fulfill_with_reg(reg)
            continue

        m = re.match(r"mov\s+([xw]\d+),\s*([xw]\d+)", s)
        if m:
            dst, src = m.group(1), m.group(2)
            if dst.startswith('w'):
                dst = 'x' + dst[1:]
            if src.startswith('w'):
                src = 'x' + src[1:]
            if src in reg_val:
                reg_val[dst] = reg_val[src]
                fulfill_with_reg(dst)
            continue

    return {reg: reg_val.get(reg) for reg in target_regs}


def combine_mov_parts(parts: Dict[int, int]) -> int:
    val = 0
    for shift, chunk in parts.items():
        val |= (chunk & 0xFFFF) << shift
    return val


def sign_extend(val: int, bits: int) -> int:
    sign_bit = 1 << (bits - 1)
    return (val & (sign_bit - 1)) - (val & sign_bit)


def decode_adrp_page_base(insn32: int, pc: int) -> Optional[int]:
    # ADRP encoding: op=1, immlo[30:29], immhi[23:5], rd[4:0]
    # target = align_down(PC, 4KB) + sign_extend((immhi:immlo) << 12)
    try:
        immlo = (insn32 >> 29) & 0x3
        immhi = (insn32 >> 5) & 0x7FFFF
        imm = (immhi << 2) | immlo  # 21 bits
        imm = sign_extend(imm, 21) << 12
        page = pc & ~0xFFF
        return (page + imm) & 0xFFFFFFFFFFFFFFFF
    except Exception:
        return None



def analyze(path: str, specs: List[TargetSpec]) -> List[CallHit]:
    with open(path, 'rb') as f:
        elf = ELFFile(f)
        mach = elf.header['e_machine']
        is_aarch64 = (mach == 'EM_AARCH64') or (mach == 183)
        if not is_aarch64:
            raise RuntimeError(f"ELF machine '{mach}' is not AArch64 (EM_AARCH64)")

        if not specs:
            return []

        lines = run_objdump_disasm(path)
        instrs = parse_objdump(lines)
        target_names = [spec.name for spec in specs]
        calls = scan_target_calls(instrs, target_names)
        if not calls:
            return []

        spec_by_name = {spec.name: spec for spec in specs}
        results: List[CallHit] = []
        for idx, func, addr, target_name in calls:
            spec = spec_by_name.get(target_name)
            if not spec:
                continue
            param_count = min(spec.max_params, MAX_REGISTER_PARAMS)
            regs = [f'x{n}' for n in range(param_count)]
            resolved = resolve_regs_before_call(elf, instrs, idx, regs) if regs else {}
            params: List[ParamHit] = []
            for p_idx, reg in enumerate(regs):
                val = resolved.get(reg) if resolved else None
                sval = read_c_string_at_va(elf, val) if isinstance(val, int) else None
                params.append(ParamHit(index=p_idx + 1, reg=reg, value=val, string=sval))
            caller = func if func != '?' else None
            results.append(CallHit(target=target_name, callsite=addr, params=params, caller=caller))
        return results




_LO12_PAT = re.compile(r"(?::lo12:)?(.+)$")


def strip_lo12_suffix(sym: str) -> str:
    m = _LO12_PAT.search(sym)
    return m.group(1) if m else sym



def parse_function_specs(values: List[str], default_limit: int) -> List[TargetSpec]:
    if default_limit < 0:
        raise SystemExit('--param-limit must be non-negative')
    specs: List[TargetSpec] = []
    for item in values:
        if not item:
            continue
        if ':' in item:
            name, limit_str = item.split(':', 1)
            if not name:
                raise SystemExit(f"Invalid --function spec: {item}")
            try:
                limit = int(limit_str) if limit_str else default_limit
            except ValueError:
                raise SystemExit(f"Invalid parameter limit in --function {item}")
        else:
            name = item
            limit = default_limit
        if limit < 0:
            raise SystemExit(f"Parameter limit must be >=0 for {name}")
        specs.append(TargetSpec(name=name, max_params=limit))
    if not specs:
        specs.append(TargetSpec(name='dlsym', max_params=default_limit))
    return specs



def lookup_symbol_va(elf: ELFFile, name: str) -> Optional[int]:
    # Try .symtab then .dynsym
    for sec_name in ('.symtab', '.dynsym'):
        sec = elf.get_section_by_name(sec_name)
        if not sec:
            continue
        try:
            for sym in sec.iter_symbols():
                if sym.name == name:
                    val = sym.entry.get('st_value', 0)
                    if val:
                        return int(val)
        except Exception:
            continue
    return None



def main():
    ap = argparse.ArgumentParser(description='Inspect selected function calls in AArch64 ELF files using objdump output')
    ap.add_argument('elf', help='Path to AArch64 ELF shared object (.so)')
    ap.add_argument('-f', '--function', action='append', dest='functions', default=[], help='Function to inspect, optionally with :N to cap parameters (e.g. dlsym:2). Repeatable.')
    ap.add_argument('--param-limit', type=int, default=2, help='Default number of parameters to resolve per function (default: 2)')
    args = ap.parse_args()

    if not os.path.isfile(args.elf):
        raise SystemExit(f"Not a file: {args.elf}")

    specs = parse_function_specs(args.functions, args.param_limit)

    try:
        results = analyze(args.elf, specs)
    except Exception as e:
        raise SystemExit(f"Error: {e}")

    print(f"File: {args.elf}")
    if not results:
        print('No matching callsites found.')
        return
    for hit in results:
        if hit.params:
            formatted = []
            for param in hit.params:
                if param.value is None:
                    val_s = '?'
                else:
                    val_s = f"0x{param.value:x}"
                if param.string is not None:
                    quoted = param.string.replace('"', '\"')
                    val_s += f" -> \"{quoted}\""
                formatted.append(f"#{param.index} {param.reg}={val_s}")
            params_str = ', '.join(formatted)
        else:
            params_str = '(no params tracked)'
        prefix = f"{hit.caller}: " if hit.caller else ''
        print(f"- {prefix}0x{hit.callsite:x}: {hit.target}({params_str})")


if __name__ == '__main__':
    main()
